# -*- coding: utf-8 -*-
# Copyright: (c) 2020-2021, Ansible Project
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
"""Concrete collection candidate management helper module."""

from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

import json
import os
import tarfile
import subprocess
from contextlib import contextmanager
from hashlib import sha256
from urllib.error import URLError
from urllib.parse import urldefrag
from shutil import rmtree
from tempfile import mkdtemp

try:
    from typing import TYPE_CHECKING
except ImportError:
    TYPE_CHECKING = False

if TYPE_CHECKING:
    from typing import (
        Any,  # FIXME: !!!111
        BinaryIO, Dict, IO,
        Iterator, List, Optional,
        Set, Tuple, Type, Union,
    )

    from ansible.galaxy.dependency_resolution.dataclasses import (
        Candidate, Requirement,
    )
    from ansible.galaxy.token import GalaxyToken

from ansible.errors import AnsibleError
from ansible.galaxy import get_collections_galaxy_meta_info
from ansible.galaxy.dependency_resolution.dataclasses import _GALAXY_YAML
from ansible.galaxy.user_agent import user_agent
from ansible.module_utils._text import to_bytes, to_native, to_text
from ansible.module_utils.common.yaml import yaml_load
from ansible.module_utils.six import raise_from
from ansible.module_utils.urls import open_url
from ansible.utils.display import Display

import yaml


display = Display()

MANIFEST_FILENAME = 'MANIFEST.json'


class ConcreteArtifactsManager:
    """Manager for on-disk collection artifacts.

    It is responsible for:
        * downloading remote collections from Galaxy-compatible servers and
          direct links to tarballs or SCM repositories
        * keeping track of local ones
        * keeping track of Galaxy API tokens for downloads from Galaxy'ish
          as well as the artifact hashes
        * caching all of above
        * retrieving the metadata out of the downloaded artifacts
    """

    def __init__(self, b_working_directory, validate_certs=True):
        # type: (bytes, bool) -> None
        """Initialize ConcreteArtifactsManager caches and costraints."""
        self._validate_certs = validate_certs  # type: bool
        self._artifact_cache = {}  # type: Dict[bytes, bytes]
        self._galaxy_artifact_cache = {}  # type: Dict[Union[Candidate, Requirement], bytes]
        self._artifact_meta_cache = {}  # type: Dict[bytes, Dict[str, Optional[Union[str, List[str], Dict[str, str]]]]]
        self._galaxy_collection_cache = {}  # type: Dict[Union[Candidate, Requirement], Tuple[str, str, GalaxyToken]]
        self._b_working_directory = b_working_directory  # type: bytes

    def get_galaxy_artifact_path(self, collection):
        # type: (Union[Candidate, Requirement]) -> bytes
        """Given a Galaxy-stored collection, return a cached path.

        If it's not yet on disk, this method downloads the artifact first.
        """
        try:
            return self._galaxy_artifact_cache[collection]
        except KeyError:
            pass

        try:
            url, sha256_hash, token = self._galaxy_collection_cache[collection]
        except KeyError as key_err:
            raise_from(
                RuntimeError(
                    'The is no known source for {coll!s}'.
                    format(coll=collection),
                ),
                key_err,
            )

        display.vvvv(
            "Fetching a collection tarball for '{collection!s}' from "
            'Ansible Galaxy'.format(collection=collection),
        )

        try:
            b_artifact_path = _download_file(
                url,
                self._b_working_directory,
                expected_hash=sha256_hash,
                validate_certs=self._validate_certs,
                token=token,
            )  # type: bytes
        except URLError as err:
            raise_from(
                AnsibleError(
                    'Failed to download collection tar '
                    "from '{coll_src!s}': {download_err!s}".
                    format(
                        coll_src=to_native(collection.src),
                        download_err=to_native(err),
                    ),
                ),
                err,
            )
        else:
            display.vvv(
                "Collection '{coll!s}' obtained from "
                'server {server!s} {url!s}'.format(
                    coll=collection, server=collection.src or 'Galaxy',
                    url=collection.src.api_server if collection.src is not None
                    else '',
                )
            )

        self._galaxy_artifact_cache[collection] = b_artifact_path
        return b_artifact_path

    def get_artifact_path(self, collection):
        # type: (Union[Candidate, Requirement]) -> bytes
        """Given a concrete collection pointer, return a cached path.

        If it's not yet on disk, this method downloads the artifact first.
        """
        try:
            return self._artifact_cache[collection.src]
        except KeyError:
            pass

        # NOTE: SCM needs to be special-cased as it may contain either
        # NOTE: one collection in its root, or a number of top-level
        # NOTE: collection directories instead.
        # NOTE: The idea is to store the SCM collection as unpacked
        # NOTE: directory structure under the temporary location and use
        # NOTE: a "virtual" collection that has pinned requirements on
        # NOTE: the directories under that SCM checkout that correspond
        # NOTE: to collections.
        # NOTE: This brings us to the idea that we need two separate
        # NOTE: virtual Requirement/Candidate types --
        # NOTE: (single) dir + (multidir) subdirs
        if collection.is_url:
            display.vvvv(
                "Collection requirement '{collection!s}' is a URL "
                'to a tar artifact'.format(collection=collection.fqcn),
            )
            try:
                b_artifact_path = _download_file(
                    collection.src,
                    self._b_working_directory,
                    expected_hash=None,  # NOTE: URLs don't support checksums
                    validate_certs=self._validate_certs,
                )
            except URLError as err:
                raise_from(
                    AnsibleError(
                        'Failed to download collection tar '
                        "from '{coll_src!s}': {download_err!s}".
                        format(
                            coll_src=to_native(collection.src),
                            download_err=to_native(err),
                        ),
                    ),
                    err,
                )
        elif collection.is_scm:
            b_artifact_path = _extract_collection_from_git(
                collection.src,
                collection.ver,
                self._b_working_directory,
            )
        elif collection.is_file or collection.is_dir or collection.is_subdirs:
            b_artifact_path = to_bytes(collection.src)
        else:
            # NOTE: This may happen `if collection.is_online_index_pointer`
            raise RuntimeError(
                'The artifact is of an unexpected type {art_type!s}'.
                format(art_type=collection.type)
            )

        self._artifact_cache[collection.src] = b_artifact_path
        return b_artifact_path

    def _get_direct_collection_namespace(self, collection):
        # type: (Candidate) -> Optional[str]
        return self.get_direct_collection_meta(collection)['namespace']  # type: ignore[return-value]

    def _get_direct_collection_name(self, collection):
        # type: (Candidate) -> Optional[str]
        return self.get_direct_collection_meta(collection)['name']  # type: ignore[return-value]

    def get_direct_collection_fqcn(self, collection):
        # type: (Candidate) -> Optional[str]
        """Extract FQCN from the given on-disk collection artifact.

        If the collection is virtual, ``None`` is returned instead
        of a string.
        """
        if collection.is_virtual:
            # NOTE: should it be something like "<virtual>"?
            return None

        return '.'.join((  # type: ignore[type-var]
            self._get_direct_collection_namespace(collection),  # type: ignore[arg-type]
            self._get_direct_collection_name(collection),
        ))

    def get_direct_collection_version(self, collection):
        # type: (Union[Candidate, Requirement]) -> str
        """Extract version from the given on-disk collection artifact."""
        return self.get_direct_collection_meta(collection)['version']  # type: ignore[return-value]

    def get_direct_collection_dependencies(self, collection):
        # type: (Union[Candidate, Requirement]) -> Dict[str, str]
        """Extract deps from the given on-disk collection artifact."""
        return self.get_direct_collection_meta(collection)['dependencies']  # type: ignore[return-value]

    def get_direct_collection_meta(self, collection):
        # type: (Union[Candidate, Requirement]) -> Dict[str, Optional[Union[str, Dict[str, str], List[str]]]]
        """Extract meta from the given on-disk collection artifact."""
        try:  # FIXME: use unique collection identifier as a cache key?
            return self._artifact_meta_cache[collection.src]
        except KeyError:
            b_artifact_path = self.get_artifact_path(collection)

        if collection.is_url or collection.is_file:
            collection_meta = _get_meta_from_tar(b_artifact_path)
        elif collection.is_dir:  # should we just build a coll instead?
            # FIXME: what if there's subdirs?
            try:
                collection_meta = _get_meta_from_dir(b_artifact_path)
            except LookupError as lookup_err:
                raise_from(
                    AnsibleError(
                        'Failed to find the collection dir deps: {err!s}'.
                        format(err=to_native(lookup_err)),
                    ),
                    lookup_err,
                )
        elif collection.is_scm:
            collection_meta = {
                'name': None,
                'namespace': None,
                'dependencies': {to_native(b_artifact_path): '*'},
                'version': '*',
            }
        elif collection.is_subdirs:
            collection_meta = {
                'name': None,
                'namespace': None,
                # NOTE: Dropping b_artifact_path since it's based on src anyway
                'dependencies': dict.fromkeys(
                    map(to_native, collection.namespace_collection_paths),
                    '*',
                ),
                'version': '*',
            }
        else:
            raise RuntimeError

        self._artifact_meta_cache[collection.src] = collection_meta
        return collection_meta

    def save_collection_source(self, collection, url, sha256_hash, token):
        # type: (Candidate, str, str, GalaxyToken) -> None
        """Store collection URL, SHA256 hash and Galaxy API token.

        This is a hook that is supposed to be called before attempting to
        download Galaxy-based collections with ``get_galaxy_artifact_path()``.
        """
        self._galaxy_collection_cache[collection] = url, sha256_hash, token

    @classmethod
    @contextmanager
    def under_tmpdir(
            cls,  # type: Type[ConcreteArtifactsManager]
            temp_dir_base,  # type: str
            validate_certs=True,  # type: bool
    ):  # type: (...) -> Iterator[ConcreteArtifactsManager]
        """Custom ConcreteArtifactsManager constructor with temp dir.

        This method returns a context manager that allocates and cleans
        up a temporary directory for caching the collection artifacts
        during the dependency resolution process.
        """
        # NOTE: Can't use `with tempfile.TemporaryDirectory:`
        # NOTE: because it's not in Python 2 stdlib.
        temp_path = mkdtemp(
            dir=to_bytes(temp_dir_base, errors='surrogate_or_strict'),
        )
        b_temp_path = to_bytes(temp_path, errors='surrogate_or_strict')
        try:
            yield cls(b_temp_path, validate_certs)
        finally:
            rmtree(b_temp_path)


def parse_scm(collection, version):
    """Extract name, version, path and subdir out of the SCM pointer."""
    if ',' in collection:
        collection, version = collection.split(',', 1)
    elif version == '*' or not version:
        version = 'HEAD'

    if collection.startswith('git+'):
        path = collection[4:]
    else:
        path = collection

    path, fragment = urldefrag(path)
    fragment = fragment.strip(os.path.sep)

    if path.endswith(os.path.sep + '.git'):
        name = path.split(os.path.sep)[-2]
    elif '://' not in path and '@' not in path:
        name = path
    else:
        name = path.split('/')[-1]
        if name.endswith('.git'):
            name = name[:-4]

    return name, version, path, fragment


def _extract_collection_from_git(repo_url, coll_ver, b_path):
    name, version, git_url, fragment = parse_scm(repo_url, coll_ver)
    b_checkout_path = mkdtemp(
        dir=b_path,
        prefix=to_bytes(name, errors='surrogate_or_strict'),
    )  # type: bytes

    # Perform a shallow clone if simply cloning HEAD
    if version == 'HEAD':
        git_clone_cmd = 'git', 'clone', '--depth=1', git_url, to_text(b_checkout_path)
    else:
        git_clone_cmd = 'git', 'clone', git_url, to_text(b_checkout_path)
    # FIXME: '--branch', version

    try:
        subprocess.check_call(git_clone_cmd)
    except subprocess.CalledProcessError as proc_err:
        raise_from(
            AnsibleError(  # should probably be LookupError
                'Failed to clone a Git repository from `{repo_url!s}`.'.
                format(repo_url=to_native(git_url)),
            ),
            proc_err,
        )

    git_switch_cmd = 'git', 'checkout', to_text(version)
    try:
        subprocess.check_call(git_switch_cmd, cwd=b_checkout_path)
    except subprocess.CalledProcessError as proc_err:
        raise_from(
            AnsibleError(  # should probably be LookupError
                'Failed to switch a cloned Git repo `{repo_url!s}` '
                'to the requested revision `{commitish!s}`.'.
                format(
                    commitish=to_native(version),
                    repo_url=to_native(git_url),
                ),
            ),
            proc_err,
        )

    return (
        os.path.join(b_checkout_path, to_bytes(fragment))
        if fragment else b_checkout_path
    )


# FIXME: use random subdirs while preserving the file names
def _download_file(url, b_path, expected_hash, validate_certs, token=None):
    # type: (str, bytes, Optional[str], bool, GalaxyToken) -> bytes
    # ^ NOTE: used in download and verify_collections ^
    b_tarball_name = to_bytes(
        url.rsplit('/', 1)[1], errors='surrogate_or_strict',
    )
    b_file_name = b_tarball_name[:-len('.tar.gz')]

    b_tarball_dir = mkdtemp(
        dir=b_path,
        prefix=b'-'.join((b_file_name, b'')),
    )  # type: bytes

    b_file_path = os.path.join(b_tarball_dir, b_tarball_name)

    display.display("Downloading %s to %s" % (url, to_text(b_tarball_dir)))
    # NOTE: Galaxy redirects downloads to S3 which rejects the request
    # NOTE: if an Authorization header is attached so don't redirect it
    resp = open_url(
        to_native(url, errors='surrogate_or_strict'),
        validate_certs=validate_certs,
        headers=None if token is None else token.headers(),
        unredirected_headers=['Authorization'], http_agent=user_agent(),
    )

    with open(b_file_path, 'wb') as download_file:  # type: BinaryIO
        actual_hash = _consume_file(resp, write_to=download_file)

    if expected_hash:
        display.vvvv(
            'Validating downloaded file hash {actual_hash!s} with '
            'expected hash {expected_hash!s}'.
            format(actual_hash=actual_hash, expected_hash=expected_hash)
        )
        if expected_hash != actual_hash:
            raise AnsibleError('Mismatch artifact hash with downloaded file')

    return b_file_path


def _consume_file(read_from, write_to=None):
    # type: (BinaryIO, BinaryIO) -> str
    bufsize = 65536
    sha256_digest = sha256()
    data = read_from.read(bufsize)
    while data:
        if write_to is not None:
            write_to.write(data)
            write_to.flush()
        sha256_digest.update(data)
        data = read_from.read(bufsize)

    return sha256_digest.hexdigest()


def _normalize_galaxy_yml_manifest(
        galaxy_yml,  # type: Dict[str, Optional[Union[str, List[str], Dict[str, str]]]]
        b_galaxy_yml_path,  # type: bytes
):
    # type: (...) -> Dict[str, Optional[Union[str, List[str], Dict[str, str]]]]
    galaxy_yml_schema = (
        get_collections_galaxy_meta_info()
    )  # type: List[Dict[str, Any]]  # FIXME: <--
    # FIXME: 👆maybe precise type: List[Dict[str, Union[bool, str, List[str]]]]

    mandatory_keys = set()
    string_keys = set()  # type: Set[str]
    list_keys = set()  # type: Set[str]
    dict_keys = set()  # type: Set[str]

    for info in galaxy_yml_schema:
        if info.get('required', False):
            mandatory_keys.add(info['key'])

        key_list_type = {
            'str': string_keys,
            'list': list_keys,
            'dict': dict_keys,
        }[info.get('type', 'str')]
        key_list_type.add(info['key'])

    all_keys = frozenset(list(mandatory_keys) + list(string_keys) + list(list_keys) + list(dict_keys))

    set_keys = set(galaxy_yml.keys())
    missing_keys = mandatory_keys.difference(set_keys)
    if missing_keys:
        raise AnsibleError("The collection galaxy.yml at '%s' is missing the following mandatory keys: %s"
                           % (to_native(b_galaxy_yml_path), ", ".join(sorted(missing_keys))))

    extra_keys = set_keys.difference(all_keys)
    if len(extra_keys) > 0:
        display.warning("Found unknown keys in collection galaxy.yml at '%s': %s"
                        % (to_text(b_galaxy_yml_path), ", ".join(extra_keys)))

    # Add the defaults if they have not been set
    for optional_string in string_keys:
        if optional_string not in galaxy_yml:
            galaxy_yml[optional_string] = None

    for optional_list in list_keys:
        list_val = galaxy_yml.get(optional_list, None)

        if list_val is None:
            galaxy_yml[optional_list] = []
        elif not isinstance(list_val, list):
            galaxy_yml[optional_list] = [list_val]  # type: ignore[list-item]

    for optional_dict in dict_keys:
        if optional_dict not in galaxy_yml:
            galaxy_yml[optional_dict] = {}

    # NOTE: `version: null` is only allowed for `galaxy.yml`
    # NOTE: and not `MANIFEST.json`. The use-case for it is collections
    # NOTE: that generate the version from Git before building a
    # NOTE: distributable tarball artifact.
    if not galaxy_yml.get('version'):
        galaxy_yml['version'] = '*'

    return galaxy_yml


def _get_meta_from_dir(
        b_path,  # type: bytes
):  # type: (...) -> Dict[str, Optional[Union[str, List[str], Dict[str, str]]]]
    try:
        return _get_meta_from_installed_dir(b_path)
    except LookupError:
        return _get_meta_from_src_dir(b_path)


def _get_meta_from_src_dir(
        b_path,  # type: bytes
):  # type: (...) -> Dict[str, Optional[Union[str, List[str], Dict[str, str]]]]
    galaxy_yml = os.path.join(b_path, _GALAXY_YAML)
    if not os.path.isfile(galaxy_yml):
        raise LookupError(
            "The collection galaxy.yml path '{path!s}' does not exist.".
            format(path=to_native(galaxy_yml))
        )

    with open(galaxy_yml, 'rb') as manifest_file_obj:
        try:
            manifest = yaml_load(manifest_file_obj)
        except yaml.error.YAMLError as yaml_err:
            raise_from(
                AnsibleError(
                    "Failed to parse the galaxy.yml at '{path!s}' with "
                    'the following error:\n{err_txt!s}'.
                    format(
                        path=to_native(galaxy_yml),
                        err_txt=to_native(yaml_err),
                    ),
                ),
                yaml_err,
            )

    return _normalize_galaxy_yml_manifest(manifest, galaxy_yml)


def _get_json_from_installed_dir(
        b_path,  # type: bytes
        filename,  # type: str
):  # type: (...) -> Dict

    b_json_filepath = os.path.join(b_path, to_bytes(filename, errors='surrogate_or_strict'))

    try:
        with open(b_json_filepath, 'rb') as manifest_fd:
            b_json_text = manifest_fd.read()
    except (IOError, OSError):
        raise LookupError(
            "The collection {manifest!s} path '{path!s}' does not exist.".
            format(
                manifest=filename,
                path=to_native(b_json_filepath),
            )
        )

    manifest_txt = to_text(b_json_text, errors='surrogate_or_strict')

    try:
        manifest = json.loads(manifest_txt)
    except ValueError:
        raise AnsibleError(
            'Collection tar file member {member!s} does not '
            'contain a valid json string.'.
            format(member=filename),
        )

    return manifest


def _get_meta_from_installed_dir(
        b_path,  # type: bytes
):  # type: (...) -> Dict[str, Optional[Union[str, List[str], Dict[str, str]]]]
    manifest = _get_json_from_installed_dir(b_path, MANIFEST_FILENAME)
    collection_info = manifest['collection_info']

    version = collection_info.get('version')
    if not version:
        raise AnsibleError(
            u'Collection metadata file `{manifest_filename!s}` at `{meta_file!s}` is expected '
            u'to have a valid SemVer version value but got {version!s}'.
            format(
                manifest_filename=MANIFEST_FILENAME,
                meta_file=to_text(b_path),
                version=to_text(repr(version)),
            ),
        )

    return collection_info


def _get_meta_from_tar(
        b_path,  # type: bytes
):  # type: (...) -> Dict[str, Optional[Union[str, List[str], Dict[str, str]]]]
    if not tarfile.is_tarfile(b_path):
        raise AnsibleError(
            "Collection artifact at '{path!s}' is not a valid tar file.".
            format(path=to_native(b_path)),
        )

    with tarfile.open(b_path, mode='r') as collection_tar:  # type: tarfile.TarFile
        try:
            member = collection_tar.getmember(MANIFEST_FILENAME)
        except KeyError:
            raise AnsibleError(
                "Collection at '{path!s}' does not contain the "
                'required file {manifest_file!s}.'.
                format(
                    path=to_native(b_path),
                    manifest_file=MANIFEST_FILENAME,
                ),
            )

        with _tarfile_extract(collection_tar, member) as (_member, member_obj):
            if member_obj is None:
                raise AnsibleError(
                    'Collection tar file does not contain '
                    'member {member!s}'.format(member=MANIFEST_FILENAME),
                )

            text_content = to_text(
                member_obj.read(),
                errors='surrogate_or_strict',
            )

            try:
                manifest = json.loads(text_content)
            except ValueError:
                raise AnsibleError(
                    'Collection tar file member {member!s} does not '
                    'contain a valid json string.'.
                    format(member=MANIFEST_FILENAME),
                )
            return manifest['collection_info']


@contextmanager
def _tarfile_extract(
        tar,  # type: tarfile.TarFile
        member,  # type: tarfile.TarInfo
):
    # type: (...) -> Iterator[Tuple[tarfile.TarInfo, Optional[IO[bytes]]]]
    tar_obj = tar.extractfile(member)
    try:
        yield member, tar_obj
    finally:
        if tar_obj is not None:
            tar_obj.close()
